import jax
import jax.numpy as jnp
import jaxsim.api as js
import resolve_robotics_uri_py
import rod

# Find the urdf file.
urdf_path = resolve_robotics_uri_py.resolve_robotics_uri(
    uri="model://ergoCubSN001/model.urdf"
)

# Build the ROD model.
rod_sdf = rod.Sdf.load(sdf=urdf_path)

# Build the model.
model = js.model.JaxSimModel.build_from_model_description(
    model_description=rod_sdf.model,
)

# Get the parent body of the collidable points.
parent_link_index_of_collidable_points = jnp.array(
    model.kin_dyn_parameters.contact_parameters.body
)

# Create the 6D forces of collidable points.
nc = len(parent_link_index_of_collidable_points)
W_f_Ci = jnp.ones(shape=(nc, 6))

# =========
# Old logic
# =========

f_old = jax.vmap(
    lambda nc: (
        jnp.vstack(jnp.equal(parent_link_index_of_collidable_points, nc).astype(int))
        * W_f_Ci
    ).sum(axis=0)
)(jnp.arange(model.number_of_links()))

# =========
# New logic
# =========

f_new = jnp.where(
    (
        parent_link_index_of_collidable_points[:, None]
        == jnp.arange(model.number_of_links())
    ).any(axis=-1, keepdims=True),
    W_f_Ci,
    jnp.zeros_like(W_f_Ci),
).sum(axis=0)